from typing import TypedDict, Annotated
from langgraph.graph import StateGraph, END
from nodes import NovelNodes, NovelState

class NovelWorkflow:
    def __init__(self):
        self.nodes = NovelNodes()
        self.graph = StateGraph(NovelState)

    def _should_refine_outline(self, state: NovelState) -> str:
        if state.get('iteration_count', 0) >= 5:
            return 'done'
        return 'refine'

    def _should_continue_chapters(self, state: NovelState) -> str:
        # 根据章节生成状态和迭代次数判断是否继续
        has_more_chapters = state.get('current_chapter_index', 0) < len(state.get('chapters_to_generate', []))
        within_limit = state.get('total_iterations', 0) < 5
        return 'continue' if has_more_chapters and within_limit else 'done'

    def compile_workflow(self):
        # 添加核心节点
        self.graph.add_node("init_outline", self.nodes.generate_initial_outline)
        self.graph.add_node("refine_outline", self.nodes.refine_outline)
        self.graph.add_node("write_chapters", self.nodes.write_chapter)
        self.graph.add_node("final_review", self.nodes.final_review)

        # 设置初始流程
        self.graph.set_entry_point ("init_outline")
        
        # 添加条件流转
        self.graph.add_conditional_edges(
            "init_outline",
            self._should_refine_outline,
            {
                "refine": "refine_outline",
                "done": "write_chapters"
            }
        )

        self.graph.add_edge("refine_outline", "init_outline")
        self.graph.add_conditional_edges(
            "write_chapters",
            self._should_continue_chapters,
            {"continue": "write_chapters", "done": "final_review"}
        )
        self.graph.add_edge("final_review", END)

        return self.graph.compile()